﻿using System;
using System.Buffers;
using System.Buffers.Binary;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using PooledGrowableBufferHelper;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using WholeSlideImageLibrary.DeepZoom;
using WholeSlideImageLibrary.Features;

namespace WholeSlideImageLibrary.Formats.OptimizedPyramidTiledImage
{
    internal class OptimizedPyramidTiledWholeSlideImage : WholeSlideImage, IDeepZoomProfileHintFeature
    {
        private FileSourceReaderCache? _source;
        private readonly long _mainBlockPointer;
        private readonly long _totalFileSize;
        private FileMainBlock _mainBlock;
        private FileTileCollectionPointerBlock[]? _levels;
        private FileMetadataBlock _metadataBlock;
        private FileAssociatedImagesBlock _associatedImagesBlock;


        private OptimizedPyramidTiledWholeSlideImage(FileSourceInfo fileSource, long mainBlockPointer)
        {
            _source = new FileSourceReaderCache(fileSource);
            _mainBlockPointer = mainBlockPointer;
            _totalFileSize = fileSource.FileSize;
        }

        private async Task<bool> InitializeAsync(FileSourceReader reader, byte[] smallBuffer, CancellationToken cancellationToken)
        {
            long totalFileSize = _totalFileSize;
            // 读取MainBlock
            int readCount = await reader.ReadAsync(_mainBlockPointer, smallBuffer.AsMemory(0, FileMainBlock.MarshalSize), cancellationToken).ConfigureAwait(false);
            if (!FileMainBlock.TryParse(smallBuffer.AsSpan(0, readCount), totalFileSize, out _mainBlock))
            {
                return false;
            }
            // 读取LevelCollectionBlock和TileCollectionBlock
            int length = _mainBlock.LevelCount * 8;
            byte[] buffer = length <= smallBuffer.Length ? smallBuffer : ArrayPool<byte>.Shared.Rent(length);
            try
            {
                readCount = await reader.ReadAsync(_mainBlock.LevelCollectionBlockPointer, buffer.AsMemory(0, length), cancellationToken).ConfigureAwait(false);
                if (readCount != length)
                {
                    return false;
                }
                // 逐级读取
                int width = _mainBlock.Width, height = _mainBlock.Height, tileSize = _mainBlock.TileSize;
                _levels = new FileTileCollectionPointerBlock[_mainBlock.LevelCount];
                for (int i = 0; i < _mainBlock.LevelCount; i++)
                {
                    long offset = BinaryPrimitives.ReadInt64LittleEndian(buffer.AsSpan(i * 8, 8));
                    if (offset < 0)
                    {
                        offset += totalFileSize;
                    }
                    _levels[i] = new FileTileCollectionPointerBlock(width, height, tileSize, offset);
                    width = (width + 1) / 2;
                    height = (height + 1) / 2;
                }
            }
            finally
            {
                if (!ReferenceEquals(buffer, smallBuffer))
                {
                    ArrayPool<byte>.Shared.Return(buffer);
                }
            }

            // 读取MetadataBlock
            if (_mainBlock.MetadataBlockPointer == 0)
            {
                _metadataBlock = new FileMetadataBlock(string.Empty, Array.Empty<FileMetadataItem>(), 0);
            }
            else
            {
                readCount = await reader.ReadAsync(_mainBlock.MetadataBlockPointer, smallBuffer.AsMemory(0, 4), cancellationToken).ConfigureAwait(false);
                if (readCount != 4)
                {
                    return false;
                }
                length = 8 + BinaryPrimitives.ReadInt32LittleEndian(smallBuffer);
                buffer = length <= smallBuffer.Length ? smallBuffer : ArrayPool<byte>.Shared.Rent(length);
                try
                {
                    readCount = await reader.ReadAsync(_mainBlock.MetadataBlockPointer + 4, buffer.AsMemory(0, length), cancellationToken).ConfigureAwait(false);
                    if (!FileMetadataBlock.TryParse(buffer.AsSpan(0, readCount), totalFileSize, out _metadataBlock))
                    {
                        return false;
                    }
                }
                finally
                {
                    if (!ReferenceEquals(buffer, smallBuffer))
                    {
                        ArrayPool<byte>.Shared.Return(buffer);
                    }
                }
            }

            // 读取AssociatedImagesBlock
            if (_mainBlock.AssociatedImagesBlockPointer == 0)
            {
                _associatedImagesBlock = new FileAssociatedImagesBlock(string.Empty, Array.Empty<FileAssociatedImageItem>(), 0);
            }
            else
            {
                readCount = await reader.ReadAsync(_mainBlock.AssociatedImagesBlockPointer, smallBuffer.AsMemory(0, 4), cancellationToken).ConfigureAwait(false);
                if (readCount != 4)
                {
                    return false;
                }
                length = 8 + BinaryPrimitives.ReadInt32LittleEndian(smallBuffer);
                buffer = length <= smallBuffer.Length ? smallBuffer : ArrayPool<byte>.Shared.Rent(length);
                try
                {
                    readCount = await reader.ReadAsync(_mainBlock.AssociatedImagesBlockPointer + 4, buffer.AsMemory(0, length), cancellationToken).ConfigureAwait(false);
                    if (!FileAssociatedImagesBlock.TryParse(buffer.AsSpan(0, readCount), totalFileSize, out _associatedImagesBlock))
                    {
                        return false;
                    }
                }
                finally
                {
                    if (!ReferenceEquals(buffer, smallBuffer))
                    {
                        ArrayPool<byte>.Shared.Return(buffer);
                    }
                }
            }

            return true;
        }

        internal static async Task<OptimizedPyramidTiledWholeSlideImage?> OpenAsync(FileSourceInfo fileSource, FileSourceReader reader, byte[] smallBuffer, long mainBlockPointer, CancellationToken cancellationToken)
        {
            Debug.Assert(smallBuffer.Length >= 256);
            var image = new OptimizedPyramidTiledWholeSlideImage(fileSource, mainBlockPointer);
            if (!await image.InitializeAsync(reader, smallBuffer, cancellationToken).ConfigureAwait(false))
            {
                return null;
            }
            return image;
        }


        #region WholeSlideImage Implementation

        public override int LevelCount => _levels!.Length;
        public override int Width => _mainBlock.Width;
        public override int Height => _mainBlock.Height;
        public int TileSize => _mainBlock.TileSize;
        public override double MicronsPerPixelX => _mainBlock.MicronsPerPixelX;
        public override double MicronsPerPixelY => _mainBlock.MicronsPerPixelY;

        public override System.Drawing.Size GetLevelSize(int level)
        {
            int levelCount = _levels!.Length;
            if ((uint)level > (uint)levelCount)
            {
                throw new ArgumentOutOfRangeException(nameof(level));
            }

            FileTileCollectionPointerBlock levelBlock = _levels[level];
            return new System.Drawing.Size(levelBlock.Width, levelBlock.Height);
        }

        #endregion

        #region DeepZoom

        public IEnumerable<DeepZoomProfile> GetDeepZoomProfileHints()
        {
            yield return new DeepZoomProfile((short)_mainBlock.TileSize, 0, EncodedImageFileFormat.Jpeg);
        }

        public ValueTask<DeepZoomTileSource?> OpenDeepZoomAsync(DeepZoomProfile profile, CancellationToken cancellationToken = default)
        {
            if (profile.TileSize == _mainBlock.TileSize && profile.Overlap == 0 && profile.FileFormat == EncodedImageFileFormat.Jpeg)
            {
                return new ValueTask<DeepZoomTileSource?>(new OptimizedPyramidTiledImageDeepZoomTileSource(this));
            }
            return default;
        }

        #endregion

        public override async ValueTask ReadRegionAsync<T>(int level, int x, int y, int width, int height, Memory<T> buffer, CancellationToken cancellationToken)
        {
            FileSourceReaderCache? cache = EnsureNotDisposed();

            int levelCount = _levels!.Length;
            if ((uint)level > (uint)levelCount)
            {
                throw new ArgumentOutOfRangeException(nameof(level));
            }
            if (Unsafe.SizeOf<T>() * buffer.Length < 4 * width * height)
            {
                throw new ArgumentException("destination is too small.", nameof(buffer));
            }

            int levelWidth = _levels[level].Width;
            int levelHeight = _levels[level].Height;

            // 修正参数
            int offsetX = Math.Min(x, levelWidth);
            int offsetY = Math.Min(y, levelHeight);
            int clippedWidth = Math.Min(width, levelWidth - offsetX);
            int clippedHeight = Math.Min(height, levelHeight - offsetY);

            // 确定需要扫描的Tiles
            int tileSize = TileSize;
            int colStart = offsetX / tileSize;
            int colEnd = (offsetX + clippedWidth + tileSize - 1) / tileSize;
            int rowStart = offsetY / tileSize;
            int rowEnd = (offsetY + clippedHeight + tileSize - 1) / tileSize;

            using var destination = Image.WrapMemory(buffer.Cast<T, Rgba32>(), width, height);

            byte[]? pooledBuffer = null;
            FileSourceReader? reader = null;
            try
            {
                reader = await cache.RentAsync(cancellationToken).ConfigureAwait(false);
                pooledBuffer = ArrayPool<byte>.Shared.Rent(FileContentRegion.MarshalSize);

                // 绘制覆盖到的Tiles到画布上
                using (var canvas = new Image<Rgba32>((colEnd - colStart) * tileSize, (rowEnd - rowStart) * tileSize))
                {
                    for (int row = rowStart; row < rowEnd; row++)
                    {
                        for (int col = colStart; col < colEnd; col++)
                        {
                            FileContentRegion tilePosition = await GetTileStreamPositionWorkspaceAsync(reader, level, col, row, pooledBuffer, cancellationToken).ConfigureAwait(false);
                            using (PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream())
                            {
                                await CopyToAsync(reader, ms, tilePosition.Offset, tilePosition.Length, cancellationToken).ConfigureAwait(false);
                                ms.Seek(0, SeekOrigin.Begin);
                                using (var tile = Image.Load(ms))
                                {
                                    canvas.Mutate(ctx => ctx.DrawImage(tile, new Point((col - colStart) * tileSize, (row - rowStart) * tileSize), opacity: 1));
                                }
                            }
                        }
                    }

                    // 从画布上返回需要的部分
                    destination.Mutate(ctx => ctx.DrawImage(canvas, new Point(colStart * tileSize - offsetX, rowStart * tileSize - offsetY), opacity: 1));
                }
            }
            finally
            {
                if (!(pooledBuffer is null))
                {
                    ArrayPool<byte>.Shared.Return(pooledBuffer);
                }
                if (!(reader is null))
                {
                    await cache.ReturnAsync(reader).ConfigureAwait(false);
                }
            }

        }


        #region Helper

        internal FileSourceReaderCache GetReaderCache()
        {
            return EnsureNotDisposed();
        }

        internal async Task<FileContentRegion> GetTileStreamPositionAsync(FileSourceReader reader, int level, int column, int row, CancellationToken cancellationToken)
        {
            int levelCount = _levels!.Length;
            if ((uint)level > (uint)levelCount)
            {
                throw new ArgumentOutOfRangeException(nameof(level));
            }

            FileTileCollectionPointerBlock levelBlock = _levels[level];
            int tileOffset = levelBlock.TileAcross * row + column;
            long offset = levelBlock.TileCollectionBlockPointer + FileContentRegion.MarshalSize * tileOffset;

            byte[] buffer = ArrayPool<byte>.Shared.Rent(FileContentRegion.MarshalSize);
            try
            {
                int bytesRead = await reader.ReadAsync(offset, buffer.AsMemory(0, FileContentRegion.MarshalSize), cancellationToken).ConfigureAwait(false);
                if (bytesRead != FileContentRegion.MarshalSize || !FileContentRegion.TryParse(buffer, _totalFileSize, out FileContentRegion tilePosition))
                {
                    return default;
                }
                return tilePosition;
            }
            finally
            {
                ArrayPool<byte>.Shared.Return(buffer);
            }
        }

        internal async Task<FileContentRegion> GetTileStreamPositionWorkspaceAsync(FileSourceReader reader, int level, int column, int row, Memory<byte> buffer, CancellationToken cancellationToken)
        {
            int levelCount = _levels!.Length;
            if ((uint)level > (uint)levelCount)
            {
                throw new ArgumentOutOfRangeException(nameof(level));
            }

            Debug.Assert(buffer.Length >= FileContentRegion.MarshalSize);

            FileTileCollectionPointerBlock levelBlock = _levels[level];
            int tileOffset = levelBlock.TileAcross * row + column;
            long offset = levelBlock.TileCollectionBlockPointer + FileContentRegion.MarshalSize * tileOffset;

            int bytesRead = await reader.ReadAsync(offset, buffer.Slice(0, FileContentRegion.MarshalSize), cancellationToken).ConfigureAwait(false);
            if (bytesRead != FileContentRegion.MarshalSize || !FileContentRegion.TryParse(buffer.Span, _totalFileSize, out FileContentRegion tilePosition))
            {
                return default;
            }
            return tilePosition;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        private FileSourceReaderCache EnsureNotDisposed()
        {
            if (_source is null)
            {
                ThrowObjectDisposedException();
            }
            return _source;
        }

        [DoesNotReturn]
        private static void ThrowObjectDisposedException()
        {
            throw new ObjectDisposedException(nameof(OptimizedPyramidTiledWholeSlideImage));
        }

        private const int DefaultCopyBufferSize = 81920;

        internal static async Task CopyToAsync(FileSourceReader source, IBufferWriter<byte> destination, long offset, int count, CancellationToken cancellationToken = default)
        {
            if (source == null)
            {
                throw new ArgumentNullException(nameof(source));
            }

            if (destination == null)
            {
                throw new ArgumentNullException(nameof(destination));
            }

            if (count < 0)
            {
                throw new ArgumentOutOfRangeException(nameof(count));
            }
            else if (count == 0)
            {
                return;
            }

            while (count > 0)
            {
                cancellationToken.ThrowIfCancellationRequested();

                int bytesToRead = Math.Min(DefaultCopyBufferSize, count);
                Memory<byte> memory = destination.GetMemory(bytesToRead);
                bytesToRead = Math.Min(bytesToRead, count);

                int bytesRead = await source.ReadAsync(offset, memory.Slice(0, bytesToRead), cancellationToken).ConfigureAwait(false);
                if (bytesRead != bytesToRead)
                {
                    throw new InvalidDataException();
                }

                destination.Advance(bytesToRead);
                offset += bytesToRead;
                count -= bytesRead;
            }
        }

        #endregion
    }
}
